#!/usr/bin/env python3
#
# Copyright 2021 Graviti. Licensed under MIT License.
#

"""The implementation of the TensorBay polygon label."""

from typing import Any, Dict, Iterable, Optional, Type, TypeVar

from tensorbay.geometry import RLE, MultiPolygon, Polygon
from tensorbay.label.basic import SubcatalogBase, _LabelBase
from tensorbay.label.supports import AttributesMixin, CategoriesMixin, IsTrackingMixin
from tensorbay.utility import ReprType, attr_base, common_loads


class PolygonSubcatalog(SubcatalogBase, IsTrackingMixin, CategoriesMixin, AttributesMixin):
    """This class defines the subcatalog for polygon type of labels.

    Arguments:
        is_tracking: A boolean value indicates whether the corresponding
            subcatalog contains tracking information.

    Attributes:
        description: The description of the entire polygon subcatalog.
        categories: All the possible categories in the corresponding dataset
            stored in a :class:`~tensorbay.utility.name.NameList`
            with the category names as keys
            and the :class:`~tensorbay.label.supports.CategoryInfo` as values.
        category_delimiter: The delimiter in category values indicating parent-child relationship.
        attributes: All the possible attributes in the corresponding dataset
            stored in a :class:`~tensorbay.utility.name.NameList`
            with the attribute names as keys
            and the :class:`~tensorbay.label.attribute.AttributeInfo` as values.
        is_tracking: Whether the Subcatalog contains tracking information.

    Examples:
        *Initialization Method 1:* Init from ``PolygonSubcatalog.loads()`` method.

        >>> catalog = {
        ...     "POLYGON": {
        ...         "isTracking": True,
        ...         "categories": [{"name": "0"}, {"name": "1"}],
        ...         "attributes": [{"name": "gender", "enum": ["male", "female"]}],
        ...     }
        ... }
        >>> PolygonSubcatalog.loads(catalog["POLYGON"])
        PolygonSubcatalog(
          (is_tracking): True,
          (categories): NameList [...],
          (attributes): NameList [...]
        )

        *Initialization Method 2:* Init an empty PolygonSubcatalog and then add the attributes.

        >>> from tensorbay.utility import NameList
        >>> from tensorbay.label import CategoryInfo, AttributeInfo
        >>> categories = NameList()
        >>> categories.append(CategoryInfo("a"))
        >>> attributes = NameList()
        >>> attributes.append(AttributeInfo("gender", enum=["female", "male"]))
        >>> polygon_subcatalog = PolygonSubcatalog()
        >>> polygon_subcatalog.is_tracking = True
        >>> polygon_subcatalog.categories = categories
        >>> polygon_subcatalog.attributes = attributes
        >>> polygon_subcatalog
        PolygonSubcatalog(
          (is_tracking): True,
          (categories): NameList [...],
          (attributes): NameList [...]
        )

    """

    def __init__(self, is_tracking: bool = False) -> None:
        SubcatalogBase.__init__(self)
        IsTrackingMixin.__init__(self, is_tracking)


class MultiPolygonSubcatalog(SubcatalogBase, IsTrackingMixin, CategoriesMixin, AttributesMixin):
    """This class defines the subcatalog for multiple polygon type of labels.

    Arguments:
        is_tracking: A boolean value indicates whether the corresponding
            subcatalog contains tracking information.

    Attributes:
        description: The description of the entire multiple polygon subcatalog.
        categories: All the possible categories in the corresponding dataset
            stored in a :class:`~tensorbay.utility.name.NameList`
            with the category names as keys
            and the :class:`~tensorbay.label.supports.CategoryInfo` as values.
        category_delimiter: The delimiter in category values indicating parent-child relationship.
        attributes: All the possible attributes in the corresponding dataset
            stored in a :class:`~tensorbay.utility.name.NameList`
            with the attribute names as keys
            and the :class:`~tensorbay.label.attribute.AttributeInfo` as values.
        is_tracking: Whether the Subcatalog contains tracking information.

    Examples:
        *Initialization Method 1:* Init from ``MultiPolygonSubcatalog.loads()`` method.

        >>> catalog = {
        ...     "MULTI_POLYGON": {
        ...         "isTracking": True,
        ...         "categories": [{"name": "0"}, {"name": "1"}],
        ...         "attributes": [{"name": "gender", "enum": ["male", "female"]}],
        ...     }
        ... }
        >>> MultiPolygonSubcatalog.loads(catalog["MULTI_POLYGON"])
        MultiPolygonSubcatalog(
          (is_tracking): True,
          (categories): NameList [...],
          (attributes): NameList [...]
        )

        *Initialization Method 2:* Init an empty MultiPolygonSubcatalog
        and then add the attributes.

        >>> from tensorbay.label import CategoryInfo, AttributeInfo
        >>> multi_polygon_subcatalog = MultiPolygonSubcatalog()
        >>> multi_polygon_subcatalog.is_tracking = True
        >>> multi_polygon_subcatalog.add_category("a")
        >>> multi_polygon_subcatalog.add_attribute("gender", enum=["female", "male"])
        >>> multi_polygon_subcatalog
        MultiPolyline2DSubcatalog(
          (is_tracking): True,
          (categories): NameList [...],
          (attributes): NameList [...]
        )

    """

    def __init__(self, is_tracking: bool = False) -> None:
        SubcatalogBase.__init__(self)
        IsTrackingMixin.__init__(self, is_tracking)


class RLESubcatalog(SubcatalogBase, IsTrackingMixin, CategoriesMixin, AttributesMixin):
    """This class defines the subcatalog for rle type of labels.

    Arguments:
        is_tracking: A boolean value indicating whether the corresponding
            subcatalog contains tracking information.

    Attributes:
        description: The description of the rle subcatalog.
        categories: All the possible categories in the corresponding dataset
            stored in a :class:`~tensorbay.utility.name.NameList`
            with the category names as keys
            and the :class:`~tensorbay.label.supports.CategoryInfo` as values.
        category_delimiter: The delimiter in category values indicating parent-child relationship.
        attributes: All the possible attributes in the corresponding dataset
            stored in a :class:`~tensorbay.utility.name.NameList`
            with the attribute names as keys
            and the :class:`~tensorbay.label.attribute.AttributeInfo` as values.
        is_tracking: Whether the Subcatalog contains tracking information.

    Examples:
        *Initialization Method 1:* Init from ``RLESubcatalog.loads()`` method.

        >>> catalog = {
        ...     "RLE": {
        ...         "isTracking": True,
        ...         "categories": [{"name": "0"}, {"name": "1"}],
        ...         "attributes": [{"name": "gender", "enum": ["male", "female"]}],
        ...     }
        ... }
        >>> RLESubcatalog.loads(catalog["RLESubcatalog"])
        RLESubcatalog(
          (is_tracking): True,
          (categories): NameList [...],
          (attributes): NameList [...]
        )

        *Initialization Method 2:* Init an empty RLESubcatalog
        and then add the attributes.

        >>> from tensorbay.label import CategoryInfo, AttributeInfo
        >>> rle_subcatalog = RLESubcatalog()
        >>> rle_subcatalog.is_tracking = True
        >>> rle_subcatalog.add_category("a")
        >>> rle_subcatalog.add_attribute("gender", enum=["female", "male"])
        >>> rle_subcatalog
        RLESubcatalog(
          (is_tracking): True,
          (categories): NameList [...],
          (attributes): NameList [...]
        )

    """

    def __init__(self, is_tracking: bool = False) -> None:
        SubcatalogBase.__init__(self)
        IsTrackingMixin.__init__(self, is_tracking)


class LabeledPolygon(_LabelBase, Polygon):
    """This class defines the concept of polygon label.

    :class:`LabeledPolygon` is the polygon type of label,
    which is often used for CV tasks such as semantic segmentation.

    Arguments:
        points: A list of 2D points representing the vertexes of the polygon.
        category: The category of the label.
        attributes: The attributs of the label.
        instance: The instance id of the label.

    Attributes:
        category: The category of the label.
        attributes: The attributes of the label.
        instance: The instance id of the label.

    Examples:
        >>> LabeledPolygon(
        ...     [(1, 2), (2, 3), (1, 3)],
        ...     category = "example",
        ...     attributes = {"key": "value"},
        ...     instance = "123",
        ... )
        LabeledPolygon [
          Vector2D(1, 2),
          Vector2D(2, 3),
          Vector2D(1, 3)
        ](
          (category): 'example',
          (attributes): {...},
          (instance): '123'
        )

    """

    _T = TypeVar("_T", bound="LabeledPolygon")

    _repr_type = ReprType.SEQUENCE
    _repr_attrs = _LabelBase._repr_attrs
    _attrs_base: Polygon = attr_base(key="polygon")

    def __init__(
        self,
        points: Optional[Iterable[Iterable[float]]] = None,
        *,
        category: Optional[str] = None,
        attributes: Optional[Dict[str, Any]] = None,
        instance: Optional[str] = None,
    ):
        Polygon.__init__(self, points)  # type: ignore[arg-type]
        _LabelBase.__init__(self, category, attributes, instance)

    @classmethod
    def loads(cls: Type[_T], contents: Dict[str, Any]) -> _T:  # type: ignore[override]
        """Loads a LabeledPolygon from a dict containing the information of the label.

        Arguments:
            contents: A dict containing the information of the polygon label.

        Returns:
            The loaded :class:`LabeledPolygon` object.

        Examples:
            >>> contents = {
            ...     "polygon": [
            ...         {"x": 1, "y": 2},
            ...         {"x": 2, "y": 3},
            ...         {"x": 1, "y": 3},
            ...     ],
            ...     "category": "example",
            ...     "attributes": {"key": "value"},
            ...     "instance": "12345",
            ... }
            >>> LabeledPolygon.loads(contents)
            LabeledPolygon [
              Vector2D(1, 2),
              Vector2D(2, 3),
              Vector2D(1, 3)
            ](
              (category): 'example',
              (attributes): {...},
              (instance): '12345'
            )

        """
        return common_loads(cls, contents)

    def dumps(self) -> Dict[str, Any]:  # type: ignore[override]
        """Dumps the current polygon label into a dict.

        Returns:
            A dict containing all the information of the polygon label.

        Examples:
            >>> labeledpolygon = LabeledPolygon(
            ...     [(1, 2), (2, 3), (1, 3)],
            ...     category = "example",
            ...     attributes = {"key": "value"},
            ...     instance = "123",
            ... )
            >>> labeledpolygon.dumps()
            {
                'category': 'example',
                'attributes': {'key': 'value'},
                'instance': '123',
                'polygon': [{'x': 1, 'y': 2}, {'x': 2, 'y': 3}, {'x': 1, 'y': 3}],
            }

        """
        return self._dumps()


class LabeledMultiPolygon(_LabelBase, MultiPolygon):  # type: ignore[misc]
    """This class defines the concept of multiple polygon label.

    :class:`LabeledMultiPolygon` is the multipolygon type of label,
    which is often used for CV tasks such as semantic segmentation.

    Arguments:
        points: A list of 2D points representing the vertices of the polygon.
        category: The category of the label.
        attributes: The attributs of the label.
        instance: The instance id of the label.

    Attributes:
        category: The category of the label.
        attributes: The attributes of the label.
        instance: The instance id of the label.

    Examples:
        >>> LabeledMultiPolygon(
        ...     [[(1.0, 2.0), (2.0, 3.0), (1.0, 3.0)], [(1.0, 4.0), (2.0, 3.0), (1.0, 8.0)]],
        ...     category = "example",
        ...     attributes = {"key": "value"},
        ...     instance = "12345",
        ... )
        LabeledMultiPolygon [
            Polygon [...],
            Polygon [...]
            ](
              (category): 'example',
              (attributes): {...},
              (instance): '12345'
            )

    """

    _T = TypeVar("_T", bound="LabeledMultiPolygon")
    _repr_type = ReprType.SEQUENCE
    _repr_attrs = _LabelBase._repr_attrs
    _attrs_base: MultiPolygon = attr_base(key="multiPolygon")

    def __init__(
        self,
        polygons: Optional[Iterable[Iterable[Iterable[float]]]] = None,
        *,
        category: Optional[str] = None,
        attributes: Optional[Dict[str, Any]] = None,
        instance: Optional[str] = None,
    ):
        MultiPolygon.__init__(self, polygons=polygons)
        _LabelBase.__init__(self, category, attributes, instance)

    @classmethod
    def loads(cls: Type[_T], contents: Dict[str, Any]) -> _T:  # type: ignore[override]
        """Loads a LabeledMultiPolygon from a list of dict containing the information of the label.

        Arguments:
            contents: A dict containing the information of the multipolygon label.

        Returns:
            The loaded :class:`LabeledMultiPolygon` object.

        Examples:
            >>> contents = {
            ...     "multiPolygon": [
            ...         [
            ...             {"x": 1.0, "y": 2.0},
            ...             {"x": 2.0, "y": 3.0},
            ...             {"x": 1.0, "y": 3.0},
            ...        ],
            ...         [{"x": 1.0, "y": 4.0}, {"x": 2.0, "y": 3.0}, {"x": 1.0, "y": 8.0}],
            ...     ],
            ...     "category": "example",
            ...     "attributes": {"key": "value"},
            ...     "instance": "12345",
            ... }
            >>> LabeledMultiPolygon.loads(contents)
            LabeledMultiPolygon [
              Polygon [...],
              Polygon [...]
            ](
              (category): 'example',
              (attributes): {...},
              (instance): '12345'
            )

        """
        return common_loads(cls, contents)

    def dumps(self) -> Dict[str, Any]:  # type: ignore[override]
        """Dumps the current multipolygon label into a dict.

        Returns:
            A dict containing all the information of the multipolygon label.

        Examples:
            >>> labeledmultipolygon = LabeledMultiPolygon(
            ...     [[(1, 2), (2, 3), (1, 3)],[(1, 2), (2, 3), (1, 3)]],
            ...     category = "example",
            ...     attributes = {"key": "value"},
            ...     instance = "123",
            ... )
            >>> labeledmultipolygon.dumps()
            {
                'category': 'example',
                'attributes': {'key': 'value'},
                'instance': '123',
                'multiPolygon': [
                    [{'x': 1, 'y': 2}, {'x': 2, 'y': 3}, {'x': 1, 'y': 3}],
                    [{"x": 1.0, "y": 4.0}, {"x": 2.0, "y": 3.0}, {"x": 1.0, "y": 8.0}]
                ]
            }

        """
        return self._dumps()


class LabeledRLE(_LabelBase, RLE):  # type: ignore[misc]
    """This class defines the concept of rle label.

    :class:`LabeledRLE` is the rle type of label,
    which is often used for CV tasks such as semantic segmentation.

    Arguments:
        rle: A rle format mask.
        category: The category of the label.
        attributes: The attributs of the label.
        instance: The instance id of the label.

    Attributes:
        category: The category of the label.
        attributes: The attributes of the label.
        instance: The instance id of the label.

    Examples:
        >>> LabeledRLE(
        ...     [272, 2, 4, 4, 2, 9],
        ...     category = "example",
        ...     attributes = {"key": "value"},
        ...     instance = "12345",
        ... )
        LabeledRLE [
          272,
          2,
          ...
        ](
            (category): 'example',
            (attributes): {...},
            (instance): '12345'
        )

    """

    _T = TypeVar("_T", bound="LabeledRLE")
    _repr_type = ReprType.SEQUENCE
    _repr_attrs = _LabelBase._repr_attrs
    _attrs_base: RLE = attr_base(key="rle")

    def __init__(
        self,
        rle: Optional[Iterable[int]] = None,
        *,
        category: Optional[str] = None,
        attributes: Optional[Dict[str, Any]] = None,
        instance: Optional[str] = None,
    ):
        RLE.__init__(self, rle)
        _LabelBase.__init__(self, category, attributes, instance)

    @classmethod
    def loads(cls: Type[_T], contents: Dict[str, Any]) -> _T:  # type: ignore[override]
        """Loads a LabeledRLE from a dict containing the information of the label.

        Arguments:
            contents: A dict containing the information of the rle label.

        Returns:
            The loaded :class:`LabeledRLE` object.

        Examples:
            >>> contents = {
            ...     "rle": [272, 2, 4, 4, 2, 9],
            ...     "category": "example",
            ...     "attributes": {"key": "value"},
            ...     "instance": "12345",
            ... }
            >>> LabeledRLE.loads(contents)
            LabeledRLE [
              272,
              2,
              ...
            ](
              (category): 'example',
              (attributes): {...},
              (instance): '12345'
            )

        """
        return common_loads(cls, contents)

    def dumps(self) -> Dict[str, Any]:  # type: ignore[override]
        """Dumps the current rle label into a dict.

        Returns:
            A dict containing all the information of the rle label.

        Examples:
            >>> labeled_rle = LabeledRLE(
            ...     [272, 2, 4, 4, 2, 9],
            ...     category = "example",
            ...     attributes = {"key": "value"},
            ...     instance = "123",
            ... )
            >>> labeled_rle.dumps()
            {
                'category': 'example',
                'attributes': {'key': 'value'},
                'instance': '123',
                'rle': [272, 2, 4, 4, 2, 9]
            }

        """
        return self._dumps()
